/*
 * This file is part of yapk (yet another (auto) process killer) driver.
 *
 * Copyright (C) 2010 by Mateusz Wojcik <mewojcik@gmail.com>
 *
 * This is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2.1 of
 * the License, or (at your option) any later version.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this software; if not, write to the Free
 * Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
 * 02110-1301 USA, or see the FSF site: http://www.fsf.org.
 */
#include <windows.h>
#include <winioctl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strsafe.h>

#include "../../common/contract.h"

#define USER_SPACE_DEVICE_NAME	"\\\\.\\Yapk_drv"
#define YAPK_SERVICE_NAME		"Yapk_drv"
#define YAPK_DRV_NAME			"yapk_drv.sys"


/**
 * \brief This function is responsible for 
 *        adding new process to killer list so that
 *        if this process will start it should be killed.
 */
BOOLEAN AddProcessToKillerList()
{
	HANDLE 					hDevice;				// Handle to virtual device.
	YapkProcessInfo			ProcessInfo;			// Holds process informations.
	UINT					NumOfBytes;				// Number of Bytes send/received.
	BOOLEAN					OperationStatus;		// Holds operation status.

	// [1] Ok. Let's start from opening device and taking
	//     handle to this file.
	hDevice = CreateFile( USER_SPACE_DEVICE_NAME,
                          GENERIC_READ | GENERIC_WRITE,
                          0,
                          NULL,
                          CREATE_ALWAYS,
                          FILE_ATTRIBUTE_NORMAL,
                          NULL );						   
	if (INVALID_HANDLE_VALUE == hDevice)
	{	/*bad case*/
		DWORD err = GetLastError();										
		printf("%s() failed at line: %u failed function name: CreateFile, error code: %d", __FUNCTION__, __LINE__, err);
		return FALSE;
	}						   

	// [2] Create contract structure.
	RtlZeroMemory(&ProcessInfo, SIZEOF_YAPK_PROCESS_INFO);

	// [3] Get process name to be killed.
	printf("Please enter process name (case sensitive) to be added to killer list: ");
	scanf("%s", ProcessInfo.ProcessName);
	printf("\n");

	// [4] Send ioctl to driver.
	OperationStatus = DeviceIoControl(	hDevice,                		// Handle to device
										IOCTL_ADD_PROCESS_TO_BE_KILLED,	// IOCTL request code
										&ProcessInfo,					// Input
										SIZEOF_YAPK_PROCESS_INFO,		// Length of input
										NULL,        					// Output
										0,  							// Length of output
										&NumOfBytes,      				// Bytes placed in buffer.
										NULL);                    		// Synchronous call

	if (TRUE != OperationStatus)
	{
		DWORD err = GetLastError();
		printf("%s() failed at line: %u failed function name: DeviceIoControl, error code: %d", __FUNCTION__, __LINE__, err);
		return FALSE;
	}

	// [5] Destruction part.
	CloseHandle(hDevice);

	return TRUE;
}

/**
 * \brief This function is responsible for 
 *        removing process from killer list so that
 *        if this process will be started on each attempt.
 */
BOOLEAN DelProcessFromKillerList()
{
	HANDLE 					hDevice;				// Handle to virtual device.
	YapkProcessInfo			ProcessInfo;			// Holds process informations.
	UINT					NumOfBytes;				// Number of Bytes send/received.
	BOOLEAN					OperationStatus;		// Holds operation status.

	// [1] Ok. Let's start from opening device and taking
	//     handle to this file.
	hDevice = CreateFile( USER_SPACE_DEVICE_NAME,
                          GENERIC_READ | GENERIC_WRITE,
                          0,
                          NULL,
                          CREATE_ALWAYS,
                          FILE_ATTRIBUTE_NORMAL,
                          NULL );						   
	if (INVALID_HANDLE_VALUE == hDevice)
	{	/*bad case*/
		DWORD err = GetLastError();										
		printf("%s() failed at line: %u failed function name: CreateFile, error code: %d", __FUNCTION__, __LINE__, err);
		return FALSE;
	}						   

	// [2] Create contract structure.
	RtlZeroMemory(&ProcessInfo, SIZEOF_YAPK_PROCESS_INFO);

	// [3] Get process name to be killed.
	printf("Please enter process name (case sensitive) to be deleted from killer list: ");
	scanf("%s", ProcessInfo.ProcessName);
	printf("\n");

	// [4] Send ioctl to driver.
	OperationStatus = DeviceIoControl(	hDevice,                		// Handle to device
										IOCTL_DEL_PROCESS_TO_BE_KILLED,	// IOCTL request code
										&ProcessInfo,					// Input
										SIZEOF_YAPK_PROCESS_INFO,		// Length of input
										NULL,        					// Output
										0,  							// Length of output
										&NumOfBytes,      				// Bytes placed in buffer.
										NULL);                    		// Synchronous call

	if (TRUE != OperationStatus)
	{
		DWORD err = GetLastError();
		printf("%s() failed at line: %u failed function name: DeviceIoControl, error code: %d", __FUNCTION__, __LINE__, err);
		return FALSE;
	}

	// [5] Destruction part.
	CloseHandle(hDevice);

	return TRUE;
}

/**
 * \brief This function is responsible for showing killer list.
 * @return TRUE in case of success, FALSE otherwise.
 */
BOOLEAN ShowKillerList()
{
	HANDLE 					hDevice;				// Handle to virtual device.
	YapkProcessInfoList		ListOfProcesses;		// List of processes.
	UINT					NumOfBytes;				// Number of Bytes send/received.
	BOOLEAN					OperationStatus;		// Holds operation status.
	ULONG					Counter = 0;
	

	// [1] Ok. Let's start from opening device and taking
	//     handle to this file.
	hDevice = CreateFile( USER_SPACE_DEVICE_NAME,
                          GENERIC_READ | GENERIC_WRITE,
                          0,
                          NULL,
                          CREATE_ALWAYS,
                          FILE_ATTRIBUTE_NORMAL,
                          NULL );						   
	if (INVALID_HANDLE_VALUE == hDevice)
	{	/*bad case*/
		DWORD err = GetLastError();										
		printf("%s() failed at line: %u failed function name: CreateFile, error code: %d", __FUNCTION__, __LINE__, err);
		return FALSE;
	}						   

	// [2] Create contract structure.
	RtlZeroMemory(&ListOfProcesses, SIZEOF_YAPK_PROCESS_INFO_LIST);


	// [4] Send ioctl to driver.
	OperationStatus = DeviceIoControl(	hDevice,                		// Handle to device
										IOCTL_SHOW_KILLER_LIST,			// IOCTL request code
										&ListOfProcesses,				// Input
										SIZEOF_YAPK_PROCESS_INFO_LIST,	// Length of input
										&ListOfProcesses,				// Output
										SIZEOF_YAPK_PROCESS_INFO_LIST,	// Length of output
										&NumOfBytes,      				// Bytes placed in buffer.
										NULL);                    		// Synchronous call

	if (TRUE != OperationStatus)
	{
		DWORD err = GetLastError();
		printf("%s() failed at line: %u failed function name: DeviceIoControl, error code: %d", __FUNCTION__, __LINE__, err);
		return FALSE;
	}

	printf("Actual killer list (%d entries):\n", ListOfProcesses.NumOfEntries);
	for (Counter = 0; Counter < ListOfProcesses.NumOfEntries; Counter++)
	{
		printf("%d) [%s]\n", Counter, ListOfProcesses.List[Counter].ProcessName);
	}

	// [5] Destruction part.
	CloseHandle(hDevice);
	
	return TRUE;
}

/**
 * \brief This function is responsible for installing yapk service
 *        in the system.
 *
 * @return TRUE in case of success, FALSE otherwise.
 */
BOOLEAN InstallService()
{
	SC_HANDLE SCManagerHandle;			// Handle to scManager.
    SC_HANDLE ServiceHandle;			// Handle to yapk service.
    TCHAR DrvPath[MAX_PATH];

	if(FALSE == GetCurrentDirectory(MAX_PATH, DrvPath))
    {
        printf("Cannot install service (%d)\n", GetLastError());
        return;
    }

	StringCbCat(DrvPath, MAX_PATH, "\\"YAPK_DRV_NAME);

	// [1] Open Service manager.
	SCManagerHandle = OpenSCManager(	NULL,                    // Local computer
										NULL,                    // ServicesActive database 
										SC_MANAGER_ALL_ACCESS);  // Full access rights 
	if (NULL == SCManagerHandle) 
    {
        printf("OpenSCManager failed (%d)\n", GetLastError());
        return;
    }

    // [2] Create the service for yapk driver.
    ServiceHandle = CreateService(	SCManagerHandle,			// SCM database 
									YAPK_SERVICE_NAME,			// Name of service 
									YAPK_SERVICE_NAME,			// Service name to display 
									SERVICE_ALL_ACCESS,			// Desired access 
									SERVICE_KERNEL_DRIVER,		// Service type 
									SERVICE_DEMAND_START,		// Start type 
									SERVICE_ERROR_NORMAL,		// Error control type 
									DrvPath,					// Path to service's binary 
									NULL,						// No load ordering group 
									NULL,						// No tag identifier 
									NULL,						// No dependencies 
									NULL,						// LocalSystem account 
									NULL);						// No password 
    if (NULL == ServiceHandle) 
    {
        printf("CreateService failed (%d)\n", GetLastError()); 
        CloseServiceHandle(SCManagerHandle);
        return;
    }
    else 
	{
		printf("Service installed successfully.\n"); 
	}


	if ( FALSE == StartService(ServiceHandle, 0, NULL) )
	{
		printf("StartService failed (%d). \n", GetLastError());
		CloseServiceHandle(ServiceHandle);
		CloseServiceHandle(SCManagerHandle);
	}
	else
	{
		printf(" Service started successfully. \n");
	}

	// [3] Destruction part.
	CloseServiceHandle(ServiceHandle); 
    CloseServiceHandle(SCManagerHandle);
	return TRUE;
}

/**
 * \brief This function is responsible for uninstalling yapk service
 *        from system.
 *
 * @return TRUE in case of success, FALSE otherwise.
 */
BOOLEAN UnInstallService()
{
	SC_HANDLE SCManagerHandle;			// Handle to scManager.
    SC_HANDLE ServiceHandle;			// Handle to yapk service.
    TCHAR szPath[MAX_PATH];
	SERVICE_STATUS  serviceStatus;

	// [1] Open Service manager.
	SCManagerHandle = OpenSCManager(	NULL,                    // Local computer
										NULL,                    // ServicesActive database 
										SC_MANAGER_ALL_ACCESS);  // Full access rights 
	if (NULL == SCManagerHandle) 
    {
        printf("OpenSCManager failed (%d)\n", GetLastError());
        return;
    }

    // [2] Open yapk service to get handle to it.
	ServiceHandle = OpenService(	SCManagerHandle,			// SCM database.
									YAPK_SERVICE_NAME,          // Name of service.
									SERVICE_ALL_ACCESS	);					// Need delete access 
    if (NULL == ServiceHandle) 
    {
        printf("OpenService failed (%d)\n", GetLastError()); 
        CloseServiceHandle(SCManagerHandle);
        return;
    }

	// [3] Try to stop service.
	ControlService(	ServiceHandle,
                    SERVICE_CONTROL_STOP,
                    &serviceStatus);


	// [4] Delete service.
	if (FALSE == DeleteService(ServiceHandle) ) 
    {
        printf("DeleteService failed (%d)\n", GetLastError()); 
    }
    else 
	{
		printf("Service deleted successfully\n"); 
	}
   
	// [5] Destruction part.
	CloseServiceHandle(ServiceHandle); 
    CloseServiceHandle(SCManagerHandle);
	return TRUE;
}

typedef enum _KillerOperation
{
	INSTALL_SERVICE					= 1,
	UNINSTALL_SERVICE,
	ADD_PROCESS_TO_KILLER_LIST,
	DEL_PROCESS_FROM_KILLER_LIST,
	SHOW_KILLER_LIST,
	EXIT
} KillerOperation;

int __cdecl main (int argc, char **argv)
{
	UINT				OperationStatus = 0;
	KillerOperation		Operation		= EXIT;
	BOOLEAN				DoNotExitLoop	= TRUE;

	do
	{
		printf("\n\nYet Another (auto) Process Killer v 1.0\n");
		printf("---------------------------------------\n");
		printf("\n");
		printf("Please choose one of the options:\n");
		printf("1. Install yapk service in the system\n");
		printf("2. Uninstall yapk service from system.\n");
		printf("3. Add process to killer list.\n");
		printf("4. Del process from killer list.\n");
		printf("5. Show killer list.\n");
		printf("6. Exit.\n");

		OperationStatus = scanf("%d", &Operation);

		switch(Operation)
		{
		case INSTALL_SERVICE:
			DoNotExitLoop = InstallService();
		break;
		case UNINSTALL_SERVICE:
			DoNotExitLoop = UnInstallService();
		break;
		case ADD_PROCESS_TO_KILLER_LIST:
			DoNotExitLoop = AddProcessToKillerList();
			break;
		case DEL_PROCESS_FROM_KILLER_LIST:
			DoNotExitLoop = DelProcessFromKillerList();
			break;
		case SHOW_KILLER_LIST:
			DoNotExitLoop = ShowKillerList();
			break;
		default:
			printf("exiting..\n");
			return 0;
		}

	}
	while(TRUE == DoNotExitLoop);

	return 0;
}

